import os
import shutil
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, QLabel, QListWidget, QShortcut
from PyQt5.QtGui import QPixmap, QKeySequence
from PyQt5.QtCore import Qt

class ImageViewer(QWidget):
    def __init__(self):
        super().__init__()

        self.image_list = []
        self.mask_list = []
        self.current_index = 0
        self.target_folder = ''
        self.initUI()

    def initUI(self):
        # Main layout
        self.layout = QHBoxLayout(self)

        # Left layout for buttons and labels
        self.left_layout = QVBoxLayout(self)

        self.select_folder_button = QPushButton('选择图片文件夹', self)
        self.select_folder_button.clicked.connect(self.select_folder)
        self.left_layout.addWidget(self.select_folder_button)

        self.folder_label = QLabel(self)
        self.left_layout.addWidget(self.folder_label)

        self.listWidget = QListWidget(self)
        self.listWidget.itemClicked.connect(self.on_item_clicked)
        self.left_layout.addWidget(self.listWidget)

        self.select_target_folder_button = QPushButton('选择目标文件夹', self)
        self.select_target_folder_button.clicked.connect(self.select_target_folder)
        self.left_layout.addWidget(self.select_target_folder_button)

        self.target_folder_label = QLabel(self)
        self.left_layout.addWidget(self.target_folder_label)

        self.prev_button = QPushButton('上一张', self)
        self.prev_button.clicked.connect(self.show_prev_image)
        self.left_layout.addWidget(self.prev_button)

        self.next_button = QPushButton('下一张', self)
        self.next_button.clicked.connect(self.show_next_image)
        self.left_layout.addWidget(self.next_button)

        self.extract_button = QPushButton('点击提取', self)
        self.extract_button.clicked.connect(self.extract_image)
        self.left_layout.addWidget(self.extract_button)

        # Right layout for image display
        self.right_layout = QHBoxLayout(self)

        self.image_label = QLabel(self)
        self.right_layout.addWidget(self.image_label)

        self.mask_label = QLabel(self)
        self.right_layout.addWidget(self.mask_label)

        # Add layouts to main layout
        self.layout.addLayout(self.left_layout)
        self.layout.addLayout(self.right_layout)

        self.setLayout(self.layout)

        # Set up shortcuts
        self.shortcut_prev = QShortcut(QKeySequence('Left'), self)
        self.shortcut_prev.activated.connect(self.show_prev_image)

        self.shortcut_next = QShortcut(QKeySequence('Right'), self)
        self.shortcut_next.activated.connect(self.show_next_image)

        self.shortcut_extract = QShortcut(QKeySequence('Space'), self)
        self.shortcut_extract.activated.connect(self.extract_image)

    def select_folder(self):
        folder_path = QFileDialog.getExistingDirectory(self, 'Select Folder')
        if folder_path:
            self.folder_label.setText(folder_path)
            all_files = sorted(os.listdir(folder_path))
            image_files = [f for f in all_files if f.endswith('_sat.jpg')]
            mask_files = [f for f in all_files if f.endswith('_mask.png')]

            # Ensure the image and mask files match up correctly
            self.image_list = [os.path.join(folder_path, f) for f in image_files if f.replace('_sat.jpg', '') in [m.replace('_mask.png', '') for m in mask_files]]
            self.mask_list = [os.path.join(folder_path, f) for f in mask_files if f.replace('_mask.png', '') in [i.replace('_sat.jpg', '') for i in image_files]]

            self.listWidget.addItems([os.path.basename(image) for image in self.image_list])
            self.show_image()

    def select_target_folder(self):
        target_folder = QFileDialog.getExistingDirectory(self, 'Select Target Folder')
        if target_folder:
            self.target_folder = target_folder
            self.target_folder_label.setText(target_folder)

    def on_item_clicked(self, item):
        self.current_index = self.listWidget.row(item)
        self.show_image()

    def show_image(self):
        if self.image_list and self.mask_list:
            image = QPixmap(self.image_list[self.current_index])
            mask = QPixmap(self.mask_list[self.current_index])
            self.image_label.setPixmap(image.scaled(600, 600, Qt.KeepAspectRatio))
            self.mask_label.setPixmap(mask.scaled(600, 600, Qt.KeepAspectRatio))
            self.listWidget.setCurrentRow(self.current_index)

    def show_prev_image(self):
        if self.current_index > 0:
            self.current_index -= 1
            self.show_image()

    def show_next_image(self):
        if self.current_index < len(self.image_list) - 1:
            self.current_index += 1
            self.show_image()

    def extract_image(self):
        if self.target_folder:
            shutil.move(self.image_list[self.current_index], self.target_folder)
            shutil.move(self.mask_list[self.current_index], self.target_folder)
            self.image_list.pop(self.current_index)
            self.mask_list.pop(self.current_index)
            self.listWidget.takeItem(self.current_index)
            if self.current_index >= len(self.image_list):
                self.current_index -= 1
            self.show_image()

if __name__ == '__main__':
    app = QApplication([])
    viewer = ImageViewer()
    viewer.show()
    app.exec_()
